import jax
import jax.numpy as jnp
import equinox as eqx
from jaxtyping import PyTree, Array, Float

class Orthogonal(eqx.Module):
    matrix: Array
    base: Array

    def __init__(self, n: int, groups: int, key: Array):
        key1, key2 = jax.random.split(key)
        lim = 1 / jnp.sqrt(n)
        self.matrix = jax.random.uniform(key1, (groups, n, n), jnp.float32, minval=-lim, maxval=lim)
        self.base = jax.random.normal(key2, (groups, n, n)) * 0.01  # Initialize with small noise


    def __call__(self):
        E = jax.vmap(get_E, in_axes=(0, 0))(self.matrix, jax.lax.stop_gradient(self.base))
        return E # (groups, n, n)
        
def get_E(matrix: Float[Array, "n n"], base: Float[Array, "n n"]) -> Array:
    """
    Computes the orthogonal matrix E from the given matrix and base.
    
    Args:
        matrix (Array): The lower triangular matrix.
        base (Array): The base vector.
        
    Returns:
        Array: The orthogonal matrix E.
    """
    A = jnp.tril(matrix, -1)
    A = A - A.T
    E = jax.lax.stop_gradient(base) @ jax.scipy.linalg.expm(A)
    return E